import torch
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import argparse
from net import CatDogClassifier

def load_model(model_path, device):
    """加载训练好的模型"""
    model = CatDogClassifier().to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()  # 设置为评估模式
    return model

def preprocess_image(image_path, img_size=224):
    """预处理图像，与训练时相同"""
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                             std=[0.229, 0.224, 0.225])
    ])
    
    image = Image.open(image_path).convert('RGB')
    original_image = image.copy()  # 保存原始图像用于显示
    input_tensor = transform(image)
    input_batch = input_tensor.unsqueeze(0)  # 添加batch维度
    
    return original_image, input_batch

def predict_image(model, image_tensor, device):
    """对图像进行预测"""
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        output = model(image_tensor)
    
    # 获取预测结果和置信度
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    confidence, predicted_class = torch.max(probabilities, 0)
    
    return predicted_class.item(), confidence.item()

def visualize_prediction(image, prediction, confidence):
    """可视化预测结果"""
    classes = ['cat', 'dog']
    
    plt.figure(figsize=(8, 6))
    plt.imshow(image)
    plt.title(f'Prediction: {classes[prediction]} (Confidence: {confidence:.2%})')
    plt.axis('off')
    
    # 添加预测结果文本
    text = f"Predicted: {classes[prediction]}\nConfidence: {confidence:.2%}"
    plt.text(10, 30, text, 
             bbox=dict(facecolor='white', alpha=0.8),
             fontsize=12)
    
    plt.show()

def main():
    parser = argparse.ArgumentParser(description='Cat/Dog Image Classifier - Single Image Inference')
    parser.add_argument('--image-path', type=str, default=r"data\test\cats\cat.4001.jpg", help='Path to the input image')
    parser.add_argument('--model-path', type=str, default=r'saved_models/fold4_epoch0_acc0.9525.pth', 
                        help='Path to the trained model')
    args = parser.parse_args()

    # 设备设置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # 加载模型
    model = load_model(args.model_path, device)
    print(f"Loaded model from: {args.model_path}")

    # 预处理图像
    original_image, input_tensor = preprocess_image(args.image_path)
    
    # 进行预测
    predicted_class, confidence = predict_image(model, input_tensor, device)
    
    # 输出结果
    classes = ['cat', 'dog']
    print("\nPrediction Result:")
    print(f"Image: {args.image_path}")
    print(f"Predicted class: {classes[predicted_class]}")
    print(f"Confidence: {confidence:.2%}")

    # 可视化结果
    visualize_prediction(original_image, predicted_class, confidence)

if __name__ == '__main__':
    main()